#include "buffer.h"
#include "common.h"
#include <string.h>
#include <time.h>

int main() {
  aclInit(NULL);
  int ret = buffer_init(DATA_SIZE);
  if (ret != 0) {
    fprintf(stderr, "buffer_init failed\n");
    return -1;
  }
  aclrtStream d2hStream, p2pStream;
  ACL_CHECK(aclrtSetDevice(FROM_DEVICE));
  ACL_CHECK(aclrtCreateStream(&d2hStream));
  ACL_CHECK(aclrtCreateStream(&p2pStream));
  aclrtEvent d2hStartEvent, d2hEndEvent, p2pStartEvent, p2pEndEvent;
  ACL_CHECK(aclrtSetDevice(FROM_DEVICE));
  ACL_CHECK(aclrtCreateEventWithFlag(&d2hStartEvent,
                                     ACL_EVENT_TIME_LINE | ACL_EVENT_SYNC));
  ACL_CHECK(aclrtCreateEventWithFlag(&d2hEndEvent,
                                     ACL_EVENT_TIME_LINE | ACL_EVENT_SYNC));
  ACL_CHECK(aclrtCreateEventWithFlag(&p2pStartEvent,
                                     ACL_EVENT_TIME_LINE | ACL_EVENT_SYNC));
  ACL_CHECK(aclrtCreateEventWithFlag(&p2pEndEvent,
                                     ACL_EVENT_TIME_LINE | ACL_EVENT_SYNC));

  int32_t canAccessPeer;
  ACL_CHECK(aclrtDeviceCanAccessPeer(&canAccessPeer, FROM_DEVICE, TO_DEVICE));
  if (canAccessPeer == 0) {
    fprintf(stderr, "Devices cannot access each other\n");
    return -1;
  }

  ACL_CHECK(aclrtSetDevice(FROM_DEVICE));
  ACL_CHECK(aclrtDeviceEnablePeerAccess(TO_DEVICE, 0));
  ACL_CHECK(aclrtSetDevice(TO_DEVICE));
  ACL_CHECK(aclrtDeviceEnablePeerAccess(FROM_DEVICE, 0));

  void *srcbuf_d_half = (void *)((uintptr_t)srcbuf_d + DATA_SIZE / 2);
  // void *dstbuf_d_half = (void *)((uintptr_t)dstbuf_d + DATA_SIZE / 2);
  void *relaybuf_h_half = (void *)((uintptr_t)relaybuf_h + DATA_SIZE / 2);
  // Warm-up iterations
  for (int i = 0; i < WARMUP_ITERATIONS; i++) {
    ACL_CHECK(aclrtSetDevice(TO_DEVICE));
    ACL_CHECK(aclrtMemset(dstbuf_d, DATA_SIZE, 0, DATA_SIZE));
    memset(relaybuf_h, 0, DATA_SIZE);
    ACL_CHECK(aclrtSetDevice(FROM_DEVICE));
    ACL_CHECK(aclrtMemcpyAsync(dstbuf_d, DATA_SIZE / 2, srcbuf_d, DATA_SIZE / 2,
                               ACL_MEMCPY_DEVICE_TO_DEVICE, p2pStream));
    ACL_CHECK(aclrtMemcpyAsync(relaybuf_h_half, DATA_SIZE / 2, srcbuf_d_half,
                               DATA_SIZE / 2, ACL_MEMCPY_DEVICE_TO_HOST,
                               d2hStream));

    ACL_CHECK(aclrtSynchronizeStream(d2hStream));
    ACL_CHECK(aclrtSynchronizeStream(p2pStream));
    if (buffer_verify_half() != 0) {
      fprintf(stderr, "data verify failed\n");
      return -1;
    }
  }

  // Measured iterations
  float totalD2hElapsedTime = 0.0f, totalP2pElapsedTime = 0.0f;
  double totalWallTime = 0.0;
  for (int i = 0; i < MEASURED_ITERATIONS; i++) {
    float d2hElapsedTime = 0.0f, p2pElapsedTime = 0.0f;
    struct timespec startWall, endWall;
    ACL_CHECK(aclrtSetDevice(TO_DEVICE));
    ACL_CHECK(aclrtMemset(dstbuf_d, DATA_SIZE, 0, DATA_SIZE));
    memset(relaybuf_h, 0, DATA_SIZE);
    ACL_CHECK(aclrtSetDevice(FROM_DEVICE));
    clock_gettime(CLOCK_MONOTONIC, &startWall);
    ACL_CHECK(aclrtRecordEvent(p2pStartEvent, p2pStream));

    ACL_CHECK(aclrtMemcpyAsync(dstbuf_d, DATA_SIZE / 2, srcbuf_d, DATA_SIZE / 2,
                               ACL_MEMCPY_DEVICE_TO_DEVICE, p2pStream));
    ACL_CHECK(aclrtRecordEvent(p2pEndEvent, p2pStream));

    ACL_CHECK(aclrtRecordEvent(d2hStartEvent, d2hStream));
    ACL_CHECK(aclrtMemcpyAsync(relaybuf_h_half, DATA_SIZE / 2, srcbuf_d_half,
                               DATA_SIZE / 2, ACL_MEMCPY_DEVICE_TO_HOST,
                               d2hStream));
    ACL_CHECK(aclrtRecordEvent(d2hEndEvent, d2hStream));

    ACL_CHECK(aclrtSynchronizeStream(d2hStream));
    ACL_CHECK(aclrtSynchronizeStream(p2pStream));
    clock_gettime(CLOCK_MONOTONIC, &endWall);

    ACL_CHECK(aclrtSetDevice(FROM_DEVICE));
    ACL_CHECK(
        aclrtEventElapsedTime(&p2pElapsedTime, p2pStartEvent, p2pEndEvent));
    ACL_CHECK(
        aclrtEventElapsedTime(&d2hElapsedTime, d2hStartEvent, d2hEndEvent));

    double wallTime = (endWall.tv_sec - startWall.tv_sec) * 1000.0 +
                      (endWall.tv_nsec - startWall.tv_nsec) / 1e6;
    totalWallTime += wallTime;
    if (buffer_verify_half() != 0) {
      fprintf(stderr, "data verify failed\n");
      return -1;
    }

    totalD2hElapsedTime += d2hElapsedTime;
    totalP2pElapsedTime += p2pElapsedTime;
  }
  totalD2hElapsedTime /= MEASURED_ITERATIONS;
  totalP2pElapsedTime /= MEASURED_ITERATIONS;
  totalWallTime /= MEASURED_ITERATIONS;
  printf("Average D2H Time: %.3f ms, Bandwidth: %.3f GB/s\n",
         totalD2hElapsedTime,
         ((float)DATA_SIZE / 2.0f) / (totalD2hElapsedTime * 1e6f));
  printf("Average P2P Time: %.3f ms, Bandwidth: %.3f GB/s\n",
         totalP2pElapsedTime,
         ((float)DATA_SIZE / 2.0f) / (totalP2pElapsedTime * 1e6f));

  float combinedTime = totalD2hElapsedTime > totalP2pElapsedTime
                           ? totalD2hElapsedTime
                           : totalP2pElapsedTime;
  printf("Combined Time: %.3f ms, Bandwidth: %.3f GB/s\n", combinedTime,
         ((float)DATA_SIZE) / (combinedTime * 1e6f));
  printf("Wall Clock Time: %.3f ms, Bandwidth: %.3f GB/s\n", totalWallTime,
         ((float)DATA_SIZE) / (totalWallTime * 1e6));
}